package com.aigc.chinamobileai.ai.config;

import com.aigc.chinamobileai.ai.domin.Conversation;
import com.aigc.chinamobileai.ai.mapper.ConversationMapper;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.time.LocalDateTime;
import java.util.List;

@Component
public class MySqlChatMemory implements ChatMemory {
    @Autowired
    private ConversationMapper conversationMapper;

    @Override
    public void add(String conversationId, List<Message> messages) {
        List<Conversation> conversations = messages.stream()
                .map(m -> {
                    Conversation c = new Conversation();
                    c.setConversationId(conversationId);
                    c.setType(m.getMessageType().getValue());
                    c.setMessage(m.getText());
                    c.setCreateTime(LocalDateTime.now());
                    return c;
                }).toList();

        conversationMapper.insert(conversations);
    }


    @Override
    public List<Message> get(String conversationId) {

        List<Conversation> conversations = conversationMapper.selectList(new LambdaQueryWrapper<Conversation>().eq(Conversation::getConversationId, conversationId));

        // 将对话记录转换为对应类型的消息对象
        return conversations.stream().map(m -> {
            Message message = null;
            //判断是助手还是用户消息,封装成对应的对象
            if (MessageType.ASSISTANT.getValue().equals(m.getType())) {
                message = new AssistantMessage(m.getMessage());
            } else if (MessageType.USER.getValue().equals(m.getType())) {
                message = new UserMessage(m.getMessage());
            }
            return message;
        }).toList();
    }

    @Override
    public void clear(String conversationId) {
        conversationMapper.delete(new LambdaQueryWrapper<Conversation>().eq(Conversation::getConversationId, conversationId));
    }
}